#include "h_mat_mul.hpp"

shmm_engine::~shmm_engine() {
  delete status_;
  for (int i = 0; i < MATRIX_TYPE::UNDEF_MATRIX_TYPE; i++) {
    delete metric_collector_[i];
    delete constant_multipler_[i];
  }
}

MATRIX_TYPE shmm_engine::parse_matrix_type(long m, long n, long k) {
  if (m == n && n == k) {
    return MATRIX_TYPE::SQUARE;
  }
  if (k >= m && k >= n) {
    return MATRIX_TYPE::MAX_K;
  }
  if (k <= m && k <= n) {
    return MATRIX_TYPE::MIN_K;
  }
  return MATRIX_TYPE::MEDIUM_K;
}

MATRIX_TYPE shmm_engine::parse_matrix_type(const cube_signature &cube) {
  ASSERT_PRINTF(cube.ndims() == 3,
                "Invalid `cube_signature` for  matrix dims\n");
  return parse_matrix_type(cube.size(0), cube.size(1), cube.size(2));
}

static inline long next_powerof2(long x) {
  int log2x = std::floor(std::log(x) / std::log(2));
  long next = 1l << log2x;
  return next < x ? next * 2l : next;
}

hmm_metric_collector *create_metric(MATRIX_TYPE type) {
  switch (type) {
  case MATRIX_TYPE::SQUARE:
    return new hmmsq_collector;
  case MATRIX_TYPE::MAX_K:
  case MATRIX_TYPE::MEDIUM_K:
    return new hmm_maxk_collector;
  case MATRIX_TYPE::MIN_K:
    return new hmm_mink_collector;
  default:
    ERROR_PRINTF("UNDEFINED MATRIX_TYPE: %d\n", type);
    return nullptr;
  }
}

cube_signature shmm_engine::get_expected_mnk(long m, long n, long k) {
  bool is_square = m == n && m == k;
  long hyper_m = is_square ? m : next_powerof2(m);
  long hyper_n = is_square ? n : next_powerof2(n);
  long hyper_k = is_square ? k : next_powerof2(k);
  return std::move(cube_signature({hyper_m, hyper_n, hyper_k}));
}

void shmm_engine::report(MATRIX_TYPE type, std::ostream &os) {
  if (type < MATRIX_TYPE::UNDEF_MATRIX_TYPE) {
    if (metric_collector_[type]) {
      metric_collector_[type]->report(status_, os);
    } else {
      INFO_PRINTF("Report unregistered engines\n");
    }
  } else {
    INFO_PRINTF("Must call cgemm() first\n");
  }
}

void shmm_engine::report(std::ostream &os) { report(status_->matrix_type()); }

void hmmsq_collector::report(const hmm_status<shmm_engine> *status,
                             std::ostream &os) {
  bool detialed = status->threads() == 1;
  // clang-format off
  os << "  * HMM Method               : " << shmm_engine::METHOD_STR[status->method()] << std::endl;
  if (final_err != 0.0) {
  os << "  * Final error bound        : " << final_err << std::endl;
  os << "    + Init error bound       : " << init_err << std::endl;
  os << "    * RotateAlign error      : x " << ra_err / init_err << std::endl;
  os << "    * ShiftCompute error     : x " << final_err / ra_err << std::endl;
  os << "      + Init error           : " << ra_err << std::endl;
  os << "      * Rotate error         : x " << sc_rot_err / ra_err << std::endl;
  os << "      * HMul error           : x " << sc_hmul_err / sc_rot_err << std::endl;
  }
  os << "  * Final noise bound        : " << final_noise << std::endl;
  os << "    * Init noise bound       : " << init_noise << std::endl;
  os << "    + RotateAlign noise      : " << ra_noise << std::endl;
  os << "    + RA correct noise       : " << ra_correct_noise << std::endl;
  os << "    + ShiftCompute noise     : " << sc_noise << std::endl;
  os << "  * Left Capacity            : " << left_cap << std::endl;
  os << "  * HMM total time           : " << get_time_unit() << " s\n";
  os << "  + Overall breakdown        :\n";
  if (no_encoding_reuse == 0.0) {
  os << "    * HMM Encoding Time      : " << get_time_hmm_encoding() << " s\n";
  os << "      * Raw data to Vec      : " << time_raw2vec << " s\n";
  os << "      * Vec to plaintext     : " << time_vec2ptxt << " s\n";
  os << "      * Reconstruct Ptxt     : " << time_re_ptxt << " s\n";
  } else {
  os << "    * HMM Encoding Time      : " << no_encoding_reuse << " s\n";
  }
  os << "    * HMM Computing Time     : " << get_time_hmm_computing() << " s\n";
  if (detialed) {
  os << "      * HAdd                 : " << time_add << " s\n";
  os << "      * Scale                : " << time_scale << " s\n";
  os << "      * HMult                : " << time_mult << " s\n";
  os << "      * Rotation             : " << time_rot << " s\n";
  }
  os << "  + Stage-by-stage breakdown :\n";
  if (no_encoding_reuse == 0.0) {
  os << "    * HMM Encoding Time      : " << get_time_hmm_encoding() << " s\n";
  os << "      - RotateAlign          : " << encoding_ra.get_time_unit() << " s\n";
  os << "        * Raw data to Vec    : " << encoding_ra.get_time_raw2vec() << " s\n";
  os << "        * Vec to plaintext   : " << encoding_ra.get_time_vec2ptxt() << " s\n";
  os << "        * Reconstruct Ptxt   : " << encoding_ra.get_time_re_ptxt() << " s\n";
  os << "      - ShiftCompute         : " << encoding_sc.get_time_unit() << " s\n";
  os << "        * Raw data to Vec    : " << encoding_sc.get_time_raw2vec() << " s\n";
  os << "        * Vec to plaintext   : " << encoding_sc.get_time_vec2ptxt() << " s\n";
  os << "        * Reconstruct Ptxt   : " << encoding_sc.get_time_re_ptxt() << " s\n";
  } else {
  os << "    * HMM Encoding Time      : " << no_encoding_reuse << " s\n";
  }
  os << "    * HMM Computing Time     : " << get_time_hmm_computing() << " s\n";
  os << "      - RotateAlign          : " << computing_ra.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_ra.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_ra.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_ra.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_ra.get_time_rot() << " s\n";
  }
  os << "      - ShiftCompute         : " << computing_sc.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_sc.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_sc.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_sc.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_sc.get_time_rot() << " s\n";
  }
  // clang-format on
}

void hmm_maxk_collector::report(const hmm_status<shmm_engine> *status,
                                std::ostream &os) {
  bool detialed = status->threads() == 1;
  // clang-format off
  os << "  * HMM Method               : " << shmm_engine::METHOD_STR[status->method()] << std::endl;
  if (final_err != 0.0) {
  os << "  * Final error bound        : " << final_err << std::endl;
  os << "    + Init error bound       : " << init_err << std::endl;
  os << "    * RotateAlign error      : x " << ra_err / init_err << std::endl;
  os << "    * ShiftCompute error     : x " << final_err / ra_err << std::endl;
  os << "      + Init error           : " << ra_err << std::endl;
  os << "      * Rotate error         : x " << sc_rot_err / ra_err << std::endl;
  os << "      * HMul error           : x " << sc_hmul_err / sc_rot_err << std::endl;
  }
  os << "  * Final noise bound        : " << final_noise << std::endl;
  os << "    * Init noise bound       : " << init_noise << std::endl;
  os << "    + RotateReplicate noise  : " << rr_noise << std::endl;
  os << "    + RotateAlign noise      : " << ra_noise << std::endl;
  os << "    + RA correct noise       : " << ra_correct_noise << std::endl;
  os << "    + ShiftCompute noise     : " << sc_noise << std::endl;
  os << "    + SumAlong noise         : " << sa_noise << std::endl;
  os << "  * Left Capacity            : " << left_cap << std::endl;
  os << "  * HMM total time           : " << get_time_unit() << " s\n";
  os << "  + Overall breakdown        :\n";
  if (no_encoding_reuse == 0.0) {
  os << "    * HMM Encoding Time      : " << get_time_hmm_encoding() << " s\n";
  os << "      * Raw data to Vec      : " << time_raw2vec << " s\n";
  os << "      * Vec to plaintext     : " << time_vec2ptxt << " s\n";
  os << "      * Reconstruct Ptxt     : " << time_re_ptxt << " s\n";
  } else {
  os << "    * HMM Encoding Time      : " << no_encoding_reuse << " s\n";
  }
  os << "    * HMM Computing Time     : " << get_time_hmm_computing() << " s\n";
  if (detialed) {
  os << "      * HAdd                 : " << time_add << " s\n";
  os << "      * Scale                : " << time_scale << " s\n";
  os << "      * HMult                : " << time_mult << " s\n";
  os << "      * Rotation             : " << time_rot << " s\n";
  }
  os << "  + Stage-by-stage breakdown :\n";
  if (no_encoding_reuse == 0.0) {
  os << "    * HMM Encoding Time      : " << get_time_hmm_encoding() << " s\n";
  os << "      - RotateAlign          : " << encoding_ra.get_time_unit() << " s\n";
  os << "        * Raw data to Vec    : " << encoding_ra.get_time_raw2vec() << " s\n";
  os << "        * Vec to plaintext   : " << encoding_ra.get_time_vec2ptxt() << " s\n";
  os << "        * Reconstruct Ptxt   : " << encoding_ra.get_time_re_ptxt() << " s\n";
  os << "      - ShiftCompute         : " << encoding_sc.get_time_unit() << " s\n";
  os << "        * Raw data to Vec    : " << encoding_sc.get_time_raw2vec() << " s\n";
  os << "        * Vec to plaintext   : " << encoding_sc.get_time_vec2ptxt() << " s\n";
  os << "        * Reconstruct Ptxt   : " << encoding_sc.get_time_re_ptxt() << " s\n";
  } else {
  os << "    * HMM Encoding Time      : " << no_encoding_reuse << " s\n";
  }
  os << "    * HMM Computing Time     : " << get_time_hmm_computing() << " s\n";
  os << "      - RotateReplicate      : " << computing_rr.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_rr.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_rr.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_rr.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_rr.get_time_rot() << " s\n";
  }
  os << "      - RotateAlign          : " << computing_ra.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_ra.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_ra.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_ra.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_ra.get_time_rot() << " s\n";
  }
  os << "      - ShiftCompute         : " << computing_sc.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_sc.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_sc.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_sc.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_sc.get_time_rot() << " s\n";
  }
  os << "      - Sum Along            : " << computing_sa.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_sa.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_sa.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_sa.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_sa.get_time_rot() << " s\n";
  }
  // clang-format on
}

void hmm_mink_collector::report(const hmm_status<shmm_engine> *status,
                                std::ostream &os) {
  bool detialed = status->threads() == 1;
  // clang-format off
  os << "  * HMM Method               : " << shmm_engine::METHOD_STR[status->method()] << std::endl;
  if (final_err != 0.0) {
  os << "  * Final error bound        : " << final_err << std::endl;
  os << "    + Init error bound       : " << init_err << std::endl;
  os << "    * RotateAlign error      : x " << ra_err / init_err << std::endl;
  os << "    * ShiftCompute error     : x " << final_err / ra_err << std::endl;
  os << "      + Init error           : " << ra_err << std::endl;
  os << "      * Rotate error         : x " << sc_rot_err / ra_err << std::endl;
  os << "      * HMul error           : x " << sc_hmul_err / sc_rot_err << std::endl;
  }
  os << "  * Final noise bound        : " << final_noise << std::endl;
  os << "    * Init noise bound       : " << init_noise << std::endl;
  os << "    + RotateReplicate noise  : " << rr_noise << std::endl;
  os << "    + RotateAlign noise      : " << ra_noise << std::endl;
  os << "    + RA correct noise       : " << ra_correct_noise << std::endl;
  os << "    + ShiftCompute noise     : " << sc_noise << std::endl;
  os << "  * Left Capacity            : " << left_cap << std::endl;
  os << "  * HMM total time           : " << get_time_unit() << " s\n";
  os << "  + Overall breakdown        :\n";
  if (no_encoding_reuse == 0.0) {
  os << "    * HMM Encoding Time      : " << get_time_hmm_encoding() << " s\n";
  os << "      * Raw data to Vec      : " << time_raw2vec << " s\n";
  os << "      * Vec to plaintext     : " << time_vec2ptxt << " s\n";
  os << "      * Reconstruct Ptxt     : " << time_re_ptxt << " s\n";
  } else {
  os << "    * HMM Encoding Time      : " << no_encoding_reuse << " s\n";
  }
  os << "    * HMM Computing Time     : " << get_time_hmm_computing() << " s\n";
  if (detialed) {
  os << "      * HAdd                 : " << time_add << " s\n";
  os << "      * Scale                : " << time_scale << " s\n";
  os << "      * HMult                : " << time_mult << " s\n";
  os << "      * Rotation             : " << time_rot << " s\n";
  }
  os << "  + Stage-by-stage breakdown :\n";
  if (no_encoding_reuse == 0.0) {
  os << "    * HMM Encoding Time      : " << get_time_hmm_encoding() << " s\n";
  os << "      - RotateAlign          : " << encoding_ra.get_time_unit() << " s\n";
  os << "        * Raw data to Vec    : " << encoding_ra.get_time_raw2vec() << " s\n";
  os << "        * Vec to plaintext   : " << encoding_ra.get_time_vec2ptxt() << " s\n";
  os << "        * Reconstruct Ptxt   : " << encoding_ra.get_time_re_ptxt() << " s\n";
  os << "      - ShiftCompute         : " << encoding_sc.get_time_unit() << " s\n";
  os << "        * Raw data to Vec    : " << encoding_sc.get_time_raw2vec() << " s\n";
  os << "        * Vec to plaintext   : " << encoding_sc.get_time_vec2ptxt() << " s\n";
  os << "        * Reconstruct Ptxt   : " << encoding_sc.get_time_re_ptxt() << " s\n";
  } else {
  os << "    * HMM Encoding Time      : " << no_encoding_reuse << " s\n";
  }
  os << "    * HMM Computing Time     : " << get_time_hmm_computing() << " s\n";
  os << "      - RotateReplicate      : " << computing_rr.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_rr.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_rr.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_rr.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_rr.get_time_rot() << " s\n";
  }
  os << "      - RotateAlign          : " << computing_ra.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_ra.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_ra.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_ra.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_ra.get_time_rot() << " s\n";
  }
  os << "      - ShiftCompute         : " << computing_sc.get_time_unit() << " s\n";
  if (detialed) {
  os << "        * HAdd               : " << computing_sc.get_time_add() << " s\n";
  os << "        * Scale              : " << computing_sc.get_time_scale() << " s\n";
  os << "        * HMult              : " << computing_sc.get_time_mult() << " s\n";
  os << "        * Rotation           : " << computing_sc.get_time_rot() << " s\n";
  }
  // clang-format on
}